import itertools
import pathlib

import torch
import torch.utils.data

import datasets
import models
import shared.checkpoint


def main():
    torch.manual_seed(1234)
    torch.set_default_device("cuda" if torch.cuda.is_available() else "cpu")

    model = models.PinnOfKortewegDeVriesEquation()
    optimizer = torch.optim.Adam(model.parameters())
    loss = torch.nn.MSELoss()

    checkpoint_manager = shared.checkpoint.CheckpointManager(pathlib.Path("./checkpoints"))
    previous_checkpoint = checkpoint_manager.last_existed()
    if previous_checkpoint:
        checkpoint = torch.load(previous_checkpoint[1])
        model.load_state_dict(checkpoint["model"])
        optimizer.load_state_dict(checkpoint["optimizer"])
        torch.set_rng_state(checkpoint["random"])
    else:
        previous_checkpoint = (0, None)

    dataset = datasets.KdVDataset()
    data_loader = torch.utils.data.DataLoader(dataset, 1000, True)

    model.train()
    for epoch in itertools.count(previous_checkpoint[0] + 1):
        for batch, (batch_t, batch_x, batch_f) in enumerate(data_loader, start=1):
            optimizer.zero_grad()
            c = (model.physics_loss() * 10 +
                 loss(model(batch_t, batch_x), batch_f) * 10)
            c.backward()
            optimizer.step()
            print(f"Epoch {epoch} Batch {batch}: {c}")
        if epoch % 10 == 0:
            torch.save({
                "model": model.state_dict(),
                "random": torch.get_rng_state(),
                "optimizer": optimizer.state_dict()
            }, checkpoint_manager.new(epoch))


if __name__ == "__main__":
    main()
